import gradio as gr
import os, asyncio
import subprocess
import pandas as pd
import numpy as np
import shutil
import tkinter as tk
import ast

from pathlib import Path
from MindAccMapper import MindAccMapper
from MindaccModel import MindAccModel
from tkinter.filedialog import askdirectory
from config import logger

mindacc_model = MindAccModel()
mapper: MindAccMapper = None

def tk_window_asksavefile(init_dir=os.getcwd()) -> str:
    window = tk.Tk()
    window.wm_attributes('-topmost', 1)
    window.withdraw()
    filename = askdirectory(initialdir=init_dir)
    return filename


async def tk_asksavefile_asy(init_dir=os.getcwd(), suffix='') -> str:
    fname = await asyncio.to_thread(tk_window_asksavefile, init_dir)
    return fname


class MindAccAnalyzer:
    # MeanAbsoluteError趋于0，RootMeanSquareError趋于0，说明测量值与真实值越近似。
    # MeanAbsoluteError趋于0，RootMeanSquareError越大，说明存在局部过大的异常值。
    # MeanAbsoluteError越大，RootMeanSquareError等于或近似MeanAbsoluteError，说明整体偏差越集中。
    # MeanAbsoluteError越大，RootMeanSquareError越大于MeanAbsoluteError，说明存在整体偏差，且整体偏差分布分散。
    analyzers = {
        "cosine_similarity": lambda a, b: np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)),
        "relative_euclidean_distance": lambda a, b: np.linalg.norm(a - b) / np.linalg.norm(a),
        "max_absolute_error": lambda a, b: np.max(np.abs(a - b)),
        "mean_absolute_error": lambda a, b: np.mean(np.abs(a - b)),
        "root_mean_square_error": lambda a, b: np.sqrt(np.mean(np.square(a - b))),
        "max_relative_error": lambda a, b: np.max(np.abs(a - b) / np.abs(a)),
        "mean_relative_error": lambda a, b: np.mean(np.abs(a - b) / np.abs(a)),
        "accumulated_relative_error": lambda a, b: np.mean(np.abs(a - b) / np.abs(a)),
        "standard_deviation": lambda a, b: np.std(np.abs(a - b) / np.abs(a)),
        "kullback_leibler_divergence": lambda a, b: np.sum(a * np.log(a / b))
    }


def process_model(model_file: gr.File, package_root_path: str, optimize: str):
    # 保存模型文件到./model目录
    if not model_file:
        return None, None
    mindacc_model.load(model_file)
    shutil.copy(model_file, Path("model"))
    model_info = f"模型路径: {mindacc_model.path}\n模型名称: {mindacc_model.name}\n输入节点: {[str(x) for x in mindacc_model.input_nodes]}\n输出节点: {[str(x) for x in mindacc_model.output_nodes]}\n输入精度: {mindacc_model.precision}"

    converted_model = mindacc_model.run_ms_converter(optimize=optimize, package_root_path=Path(package_root_path))
    # 防止permission denied
    if (Path('model') / converted_model.name).exists():
        os.chmod(Path('model') / converted_model.name, 0o777)
    shutil.copy(converted_model, Path('model'))
    converted_model = str(converted_model)

    return model_info, gr.File(label="转换后的模型", value=converted_model, visible=True)


def process_input(input_file: gr.File):
    mindacc_model.input_load(input_file)

def get_mslite_output_shape(name):
    if name is None or len(name) == 0:
        return None
    info = MindAccMapper.get_ms_bin_info(Path(name))
    if info is None:
        return None
    else:
        return info.get('shape')

def get_onnx_output_shape(name):
    global mapper
    if mapper is None:
        mapper = MindAccMapper(onnx_model=mindacc_model.path, onnx_dump_file=mindacc_model.onnx_output_path,
                  ms_dump_dir=mindacc_model.ms_output_path)
    return str(mapper.onnx_dump[name].shape)

def random_input(seed):
    # 生成随机输入
    mindacc_model.input_generate(seed=seed)
    onnx_shape = mindacc_model.input_nodes[0].shape
    # # 4维输入时 nchw to nhwc
    ms_shape = onnx_shape if len(onnx_shape) != 4 else (onnx_shape[0], onnx_shape[2], onnx_shape[3], onnx_shape[1])
    return "input/mslite_input.bin", ms_shape


def run_infer(device: str, enable_fp16_box: bool, input_shape: str, msLite_path: str):
    # 运行推理
    ms_bench_dir = mindacc_model.run_ms_dump(device, enable_fp16_box, input_shape, Path(mindacc_model.ms_model_path), Path(msLite_path))
    mindacc_model.run_onnx_dump(dump_path=ms_bench_dir / "onnx_dump.npz")
    return f"mslite模型输出文件夹: {str(ms_bench_dir)}, onnx模型输出文件: {str(ms_bench_dir)}"


def run_map():
    global mapper
    # 构造映射
    mapper = MindAccMapper(onnx_model=mindacc_model.path, onnx_dump_file=mindacc_model.onnx_output_path,
                           ms_dump_dir=mindacc_model.ms_output_path)
    mindacc_model.compare_map = mapper.simple_map()
    # 构造pd.DataFrame用于显示
    model_map_dataframe = pd.DataFrame(mindacc_model.compare_map.items(),
                                       columns=["MSlite输出节点", "ONNX输出节点"])
    # 在DataFrame中添加未映射的节点
    maped_count, all_count, map_rate, unmap_list, map_list = mapper.get_map_result()
    unmap_df = pd.DataFrame({"MSlite输出节点": unmap_list, "ONNX输出节点": "未映射"})
    model_map_dataframe = pd.concat([model_map_dataframe, unmap_df], ignore_index=True)
    return f"映射成功率: {maped_count}/{all_count}={map_rate}", model_map_dataframe


def run_compare(analyzers, model_mapping, progress=gr.Progress()):
    nhwc2nchw = lambda x: np.transpose(x, (0, 3, 1, 2)) if len(x.shape) == 4 else x
    # 更新映射
    old = mindacc_model.compare_map
    mindacc_model.compare_map = {row["MSlite输出节点"]: row["ONNX输出节点"] for row in
                                 model_mapping.to_dict(orient='records') if
                                 row["ONNX输出节点"] != "未映射" and row["ONNX输出节点"] != ""}
    # 运行对比
    compare_result = pd.DataFrame()
    onnx_dump = np.load(mindacc_model.onnx_output_path)
    for i in progress.tqdm(mindacc_model.compare_map):
        i_path = mindacc_model.ms_output_path / i
        i_info = MindAccMapper.get_ms_bin_info(i_path)
        a_shape = i_info['shape']
        a_dtype = i_info['data_type']
        a = np.fromfile(i_path, dtype=a_dtype).reshape(a_shape)
        b = onnx_dump[mindacc_model.compare_map[i]]
        # a, b输入维度不一致时，a转换为nchw
        if a.shape != b.shape:
            a = nhwc2nchw(a)
        for analyzer in analyzers:
            # 设置节点名
            compare_result.loc[i, "节点"] = i
            try:
                compare_result.loc[i, analyzer] = MindAccAnalyzer.analyzers[analyzer](a.flatten(), b.flatten())
            except Exception as e:
                logger.warning(f"Can't run analyzer {analyzer} for node {i}, error: {e}")

    mindacc_model.compare_result = compare_result
    # 保存对比结果
    date_time = pd.Timestamp.now().strftime("%Y-%m-%d_%H-%M-%S")
    file_path = mindacc_model.ms_output_path / f"compare_result_{date_time}.csv"
    mindacc_model.compare_result.to_csv(file_path)
    return compare_result, gr.DownloadButton(label=f"下载对比结果 {file_path.name}", value=file_path, visible=True)


def transpose_output(output, isTranspose, isReshape, reshape):
    temp = output
    if isTranspose:
        temp = temp.transpose()
    if isReshape and reshape is not None:
        temp = temp.reshape(reshape)
    return temp


def run_single_compare(analyzers, msLite_name, onnx_name, transpose_msLite, reshape_msLite, new_shape_msLite,
                       transpose_onnx, reshape_onnx, new_shape_onnx):
    # 获取向量
    info = MindAccMapper.get_ms_bin_info(Path(msLite_name))
    shape = info['shape']
    dtype = info['data_type']
    ms_dump = np.fromfile(mindacc_model.ms_output_path / msLite_name, dtype=dtype).reshape(shape)
    onnx_dump = np.load(mindacc_model.onnx_output_path)[onnx_name]

    # 更新向量
    ms_dump = transpose_output(ms_dump, transpose_msLite, reshape_msLite, new_shape_msLite)
    onnx_dump = transpose_output(onnx_dump, transpose_onnx, reshape_onnx, new_shape_onnx)
    # 运行对比
    compare_result = pd.DataFrame()
    compare_result.loc[msLite_name, "节点"] = msLite_name
    for analyzer in analyzers:
        # 设置节点名
        try:
            compare_result.loc[msLite_name, analyzer] = MindAccAnalyzer.analyzers[analyzer](ms_dump.flatten(), onnx_dump.flatten())
        except Exception as e:
            logger.warning(f"Can't run analyzer {analyzer} for node , error: {e}")
    return compare_result


def get_hdc_devices():
    try:
        result = subprocess.run("hdc list targets", text=True, capture_output=True).stdout.strip()
    except:
        result = "[Empty]"
    if result == "[Empty]":
        devices = []
    else:
        devices = result.stdout.strip().split("\n")
    devices.append('local')
    return devices


def create_gradio_interface():
    with gr.Blocks(theme=gr.themes.Soft()) as mindacc:
        gr.HTML("""
            <div style="text-align: center; margin-bottom: 200px;">
                <h1 style="font-size: 3em;">MindAcc</h1>
                <p style="font-size: 1.5em;">MindAcc是一个用于对比MindSpore Lite与ONNX模型推理结果的工具，</p>
                <p style="font-size: 1.5em;">支持模型转换、随机输入生成、推理运行、对比分析等功能。</p>
            </div>
        """)
        with gr.Accordion(label="使用说明"):
            gr.HTML("""
                <div>
                <p style="font-size: 12px;">1. 上传ONNX模型文件，将自动转换为MSlite模型并读取信息</p>
                <p style="font-size: 12px;">2. 上传MSlite输入数据或点击随机生成输入按钮</p>
                <p style="font-size: 12px;">3. 点击运行推理按钮，将使用benchmark工具获取MSlite输出，使用onnxdumper获取ONNX输出</p>
                <p style="font-size: 12px;">4. 点击生成映射按钮，将自动匹配MSlite输出节点与ONNX输出节点，若需修改映射可在节点映射表中直接修改</p>
                <p style="font-size: 12px;">5. 点击运行对比按钮，选择对比分析指标，获取对比结果</p>
            </div>
            """)
        with gr.Row():
            with gr.Column(scale=1):
                random_button = gr.Button("随机生成输入", variant="primary")
                devices_list_button = gr.Dropdown(label="hdc devices", choices=get_hdc_devices())

                norm_input_text = gr.Text(label='mslite directory')
                norm_input_btn = gr.Button('Select mslite directory')
                norm_input_btn.click(tk_asksavefile_asy, inputs=[], outputs=[norm_input_text])

            with gr.Column(scale=1):
                random_seed_input = gr.Number(label="随机数种子")
                optimize_input = gr.Dropdown(label="convert-lite优化选项",
                                             choices=["none", "general", "gpu_oriented", "ascend_oriented"],
                                             value="none")
                model_info_output = gr.Textbox(label="模型信息")
            with gr.Column(scale=1):
                with gr.Group():
                    model_input = gr.File(label="上传模型")
                    converted_model_output = gr.File(label="转换后的模型", visible=False)
            with gr.Column(scale=1):
                with gr.Group():
                    data_input = gr.File(label="上传 MSlite 输入数据, 或使用随机生成")
                    ms_shape_input = gr.Textbox(label="MSlite 输入形状")

        with gr.Row():
            with gr.Column(scale=1):
                enable_fp16_box = gr.Checkbox(label='enable fp16')
                run_infer_button = gr.Button("运行推理", variant="primary", scale=1)
            infer_output = gr.Textbox(label="推理结果", scale=2)

        with gr.Row():
            run_map_button = gr.Button("生成映射", variant="primary", scale=1)
            mapping_outline = gr.Textbox(label="映射概览", scale=2)
        with gr.Row():
            model_mapping = gr.DataFrame(label="节点映射", headers=["MSlite输出节点", "ONNX输出节点"],
                                         datatype=["str", "str"])

        with gr.Row():
            run_compare_button = gr.Button("运行对比", variant="primary", scale=1)
            analysis_tool = gr.CheckboxGroup([x for x in MindAccAnalyzer.analyzers.keys()], label="对比分析指标",
                                             scale=2)

        compare_output = gr.DataFrame(label="对比结果", headers=["节点", *MindAccAnalyzer.analyzers.keys()])
        with gr.Row():
            run_draw_button = gr.Button("绘制统计图", variant="primary")
            dowload_compare_button = gr.DownloadButton("下载对比结果", variant="primary")

        # 渲染柱形图
        @gr.render(triggers=[run_draw_button.click])
        def create_plot():
            dataframe = mindacc_model.compare_result
            if '节点' not in dataframe.columns or len(dataframe['节点']) == 0:
                print("No data")
                gr.Markdown("## 等待对比结果")
            else:
                # 每一个指标绘制一个柱形图
                # 获取dataframe的列名
                analyzers = dataframe.columns[1:]
                for analyzer in analyzers:
                    # 取出指标数据
                    data = pd.DataFrame({"节点": dataframe['节点'], analyzer: dataframe[analyzer]})
                    if len(data) == 0:
                        continue
                    # 绘制柱形图
                    gr.BarPlot(data, x="节点", y=analyzer)
        gr.Markdown('## 输出比较')
        with gr.Row():
            with gr.Column():
                msLite_name = gr.Textbox(scale=0.5, label="MSLite name")
                msLite_shape = gr.Text(label="MSLite shape", lines=1, interactive=False)
            with gr.Column():
                onnx_name = gr.Textbox(scale=0.5, label="ONNX name")
                onnx_shape = gr.Text(label="ONNX shape")
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### Operation on MsLite")
                    transpose_msLite = gr.Checkbox(label="Transpose MSLite")
                    reshape_msLite = gr.Checkbox(label="Reshape MSLite")
                    new_shape_msLite = gr.Textbox(label="New Shape for MSLite", placeholder="Leave empty if no reshape")
                with gr.Column():
                    gr.Markdown("### Operation on ONNX")
                    transpose_onnx = gr.Checkbox(label="Transpose ONNX")
                    reshape_onnx = gr.Checkbox(label="Reshape ONNX")
                    new_shape_onnx = gr.Textbox(label="New Shape for ONNX",
                                                placeholder="Leave empty if no reshape")

        with gr.Row():
            run_single_compare_button = gr.Button("运行对比", variant="primary", scale=1)
            single_analysis_tool = gr.CheckboxGroup([x for x in MindAccAnalyzer.analyzers.keys()], label="对比分析指标",
                                             scale=2)
        single_compare_output = gr.DataFrame(label="对比结果", headers=["节点", *MindAccAnalyzer.analyzers.keys()])


        model_input.upload(
            fn=process_model,
            inputs=[model_input, norm_input_text, optimize_input],
            outputs=[model_info_output, converted_model_output],
        )
        data_input.upload(
            fn=process_input,
            inputs=[data_input],
            outputs=[],
        )

        random_button.click(fn=random_input, inputs=random_seed_input, outputs=[data_input, ms_shape_input])
        run_infer_button.click(fn=run_infer, inputs=[devices_list_button, enable_fp16_box, ms_shape_input, norm_input_text], outputs=[infer_output])
        run_map_button.click(fn=run_map, inputs=None, outputs=[mapping_outline, model_mapping])
        run_compare_button.click(fn=run_compare, inputs=[analysis_tool, model_mapping],
                                 outputs=[compare_output, dowload_compare_button])
        msLite_name.change(fn=get_mslite_output_shape, inputs=[msLite_name], outputs=msLite_shape)
        onnx_name.change(fn=get_onnx_output_shape, inputs=[onnx_name], outputs=onnx_shape)
        run_single_compare_button.click(fn=run_single_compare, inputs=[single_analysis_tool, msLite_name, onnx_name, transpose_msLite, reshape_msLite, new_shape_msLite, transpose_onnx, reshape_onnx, new_shape_onnx], outputs=single_compare_output)

    mindacc.launch()
